package anonymouschat.client;

import java.io.IOException;
import java.io.ObjectOutputStream;
import java.net.Socket;
import java.security.NoSuchAlgorithmException;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;

import javax.crypto.SecretKey;
import javax.swing.JOptionPane;
import javax.swing.JPasswordField;
import javax.swing.JProgressBar;

import anonymouschat.client.frames.NickListFrame;
import anonymouschat.client.network.ClientInConnectionThread;
import anonymouschat.exceptions.DecryptionException;
import anonymouschat.messages.AESkeyMessage;
import anonymouschat.messages.EncryptedMessageWrapperMessage;
import anonymouschat.messages.Message;
import anonymouschat.messages.clientrequests.KeyRequest;
import anonymouschat.messages.clientrequests.UserLoginRequest;
import anonymouschat.messages.securitymessages.RSAEncryptedMessage;
import anonymouschat.messages.securitymessages.TextMessageWithSig;
import anonymouschat.security.KeySalt;
import anonymouschat.security.Keys;
import anonymouschat.security.NameSaltPairs;
import anonymouschat.security.SecurityUtils;
import anonymouschat.service.ServerEntry;
import anonymouschat.utils.Constants;
import anonymouschat.utils.ListOfServers;
import anonymouschat.utils.LocalFileOps;
import anonymouschat.utils.Print;

/**
 * Client program.
 * 
 * @author hb
 * 
 */
public class Client {

	// stream to write messages to the initail server in route
	private ObjectOutputStream out;

	// connection to initial server (first server in the route)
	private Socket clientSocket;

	// Full Identity (Private keys for the user running this client)
	private Identity identity;

	// Public keys and user name of the user running this client
	private PublicIdentity publicIdentity;

	// GUI for the nick name list
	private NickListFrame nickListFrame;

	// substance of the nick name list (pairs of user names to user assigned
	// nicknames)
	private NickNameList nickList;

	private ListOfServers listOfAvailableServers;

	// map of other clients' public keys
	private HashMap<String, PublicIdentity> publicIdentityMap;

	private StringBuffer chain;

	/**
	 * List of servers used by the client to send/receive messages, first is the
	 * initail server Last is the publication server of this client. Fully set
	 * up when establishing anonymous route
	 */
	public static LinkedList<ServerEntry> serverChain = new LinkedList<ServerEntry>();

	/**
	 * Starting client
	 * 
	 * @throws ClassNotFoundException
	 * @throws IOException
	 * @throws NoSuchAlgorithmException
	 */
	public Client() throws ClassNotFoundException, NoSuchAlgorithmException {
		Print.info("Starting Client. . .");
		publicIdentityMap = new HashMap<String, PublicIdentity>();
		identity = selectIdentity();
		publicIdentity = new PublicIdentity(identity);
		listOfAvailableServers = new ListOfServers(this);
		try {
			establishAnonymousRouting();
		} catch (Exception e) {
			Print.error(
					"ERROR encountered while trying to establish annonymous routing: "
							+ e.getLocalizedMessage(), e);
		}
		sendUserLoginRequest();
	}

	/**
	 * Send the initial login request to a publication server
	 * 
	 * @param out
	 *            an output stream to write the message
	 */
	private void sendUserLoginRequest() {
		UserLoginRequest req = new UserLoginRequest(getIdentity(),
				serverChain.getLast().serverName);
		try {
			Print.debug(getIdentity().getUserName()
					+ " is about to send a request to login");
			sendMessage(req, true);
		} catch (IOException e) {
			Print.error("Unable to send UserLoginRequest to server", e);
		}
	}

	/**
	 * when client wants to send a chat message, it check for the client public
	 * key if it does not know it yet, it will request it first then it will
	 * wrap the message and send it
	 * 
	 * @param msg
	 */
	public void sendChatMessage(TextMessageWithSig msg) {
		// check for public key
		// if does not exist, ask for it
		String otherUser = msg.getTextMessage().getDstUser();
		if (!publicIdentityMap.containsKey(otherUser)
				|| publicIdentityMap.get(otherUser) == null) {
			// need to request keys from the server
			requestKey(otherUser);
		}
		if (waitForKeys(msg)) {
			if (publicIdentityMap.get(msg.getTextMessage().getDstUser()) == null) {
				// user is not online
				nickListFrame.showTxtOnUserWindow(msg.getTextMessage()
						.getDstUser(), msg.getTextMessage().getDstUser()
						+ " is not online.");
				Print.debug("user" + msg.getTextMessage().getDstUser()
						+ " is not online.");
			} else {
				Print.debug("Get public key of client: "
						+ msg.getTextMessage().getDstUser());
				Print.debug(publicIdentityMap.get(
						msg.getTextMessage().getDstUser()).toString());

				PublicIdentity publicIdentity = publicIdentityMap.get(msg
						.getTextMessage().getDstUser());
				if (publicIdentity == null) {
					// user is not online.
					Print.debug("user is not online.");
					return;
				}
				// send encrypted message
				RSAEncryptedMessage emsg = SecurityUtils.encRSA(msg,
						publicIdentity.getPublicRSAKey(), msg.getTextMessage()
								.getDstUser(), this.getPublicIdentity());
				try {
					Print.debug(getIdentity().getUserName()
							+ " is about to send an encrypted message to user "
							+ msg.getTextMessage().getDstUser());
					sendMessage(emsg);
				} catch (IOException e) {
					Print.error("Unable to send encrypted message to server", e);
				}
			}
		} else {
			nickListFrame
					.showTxtOnUserWindow(msg.getTextMessage().getDstUser(),
							"Request public keys for user "
									+ msg.getTextMessage().getDstUser()
									+ " timed out.");
			Print.error("Cannot find keys. Request for public keys timed out...");
		}
	}

	/**
	 * wait for keys from the server, or until the request times out.
	 * 
	 * @param msg
	 * @return true if we have the keys for the client false if we don't have
	 *         the keys after a period of time
	 */
	public boolean waitForKeys(TextMessageWithSig msg) {
		int waitCount = 0;
		boolean haveKeys = false;
		int waitLimit = 1000;
		while (!publicIdentityMap
				.containsKey(msg.getTextMessage().getDstUser())
				&& waitCount < waitLimit) {
			// wait until we get the public keys are returned from the server
			try {
				Thread.sleep(100);
			} catch (InterruptedException e) {
				Print.error("While waiting for public keys...", e);
			}
			waitCount++;
		}
		if (waitCount < waitLimit) {
			// the request did not time out
			haveKeys = true;
		}
		return haveKeys;
	}

	/**
	 * Request the public keys for a user
	 * 
	 * @param user
	 *            User who's keys we want from the server.
	 */
	public void requestKey(String user) {
		KeyRequest reqPK = new KeyRequest(getIdentity().getUserName(), user);
		try {
			Print.debug(getIdentity().getUserName()
					+ " is about to send a request for user " + user
					+ "'s public key");
			sendMessage(reqPK);
		} catch (IOException e) {
			Print.error("Unable to send PublicKeyRequest to server", e);
		}
	}

	private LinkedList<ServerEntry> pickServerChain(ServerEntry avoidAsInitial)
			throws Exception {
		ServerEntry randomServer;
		LinkedList<ServerEntry> chain = new LinkedList<ServerEntry>();

		// Choose servers to create chain
		for (int i = 0; i < Constants.numServerHops; i++) {
			do
				randomServer = listOfAvailableServers.getRandomServer(chain);
			while (i == 0 && randomServer == avoidAsInitial);
			chain.add(randomServer);
		}
		return chain;
	}

	/**
	 * /** Gets a list of random servers to use for routing
	 * 
	 * @throws Exception
	 */
	private void establishAnonymousRouting() throws Exception {
		Client.serverChain = pickServerChain(null);

		chain = new StringBuffer();
		for (int i = 0; i < Client.serverChain.size(); i++) {
			chain = chain.append(Client.serverChain.get(i).serverName);
			if (i != Client.serverChain.size() - 1)
				chain = chain.append("->");
		}
		Print.info("Server chain: " + chain.toString());

		// Establish connection to first server so we can send objects
		ServerEntry initialServer = Client.serverChain.get(0);
		clientSocket = new Socket(initialServer.serverIP,
				initialServer.serverPort);
		clientSocket.setKeepAlive(true);
		Print.debug("Connected to server " + initialServer.serverIP + ":"
				+ initialServer.serverPort);

		ClientInConnectionThread t = new ClientInConnectionThread(this,
				clientSocket);
		t.start();

		out = t.getOutputStream();
		Print.info("Established connection to " + initialServer.serverName);

		ServerEntry currentServer;
		// Generate symmetric keys and send to all servers
		for (int i = 0; i < Constants.numServerHops; i++) {
			currentServer = Client.serverChain.get(i);
			AESkeyMessage k = new AESkeyMessage();

			k.sharedAESKey = anonymouschat.security.Keys.generateAESkey();
			currentServer.setSymmetricKey(k.sharedAESKey);
			EncryptedMessageWrapperMessage enc = new EncryptedMessageWrapperMessage(
					listOfAvailableServers.getByName(currentServer.serverName)
							.getAsymmetricCipherEngine(), k,
					Constants.ASYMMETRIC_ENCRYPTION);
			sendMessage(enc, currentServer, true, true);
		}
	}

	/**
	 * Sends the message to the publication server of this client.
	 * 
	 * @param msg
	 *            message to send
	 * @throws IOException
	 *             Thrown if can't write object to stream
	 */
	public void sendMessage(Message msg) throws IOException {
		sendMessage(msg, Client.serverChain.get(Constants.numServerHops - 1),
				false, true); // Send to the published server
	}

	public void sendMessage(Message msg, boolean useReturnChain)
			throws IOException {
		sendMessage(msg, Client.serverChain.get(Constants.numServerHops - 1),
				false, useReturnChain);
	}

	/**
	 * Sends message to a destination server in this user's route
	 * 
	 * @param msgin
	 *            message to send
	 * @param destination
	 *            destination server that the message is intended for
	 * @param sendingAESkey
	 *            true if sending the server the key we want it to use for onion
	 *            wrapping on the way back
	 * @throws IOException
	 */
	public void sendMessage(Message msgin, ServerEntry destination,
			boolean sendingAESkey, boolean useReturnChain) throws IOException {
		try {
			LinkedList<ServerEntry> chain;
			if (useReturnChain)
				chain = Client.serverChain;
			else
				chain = pickServerChain(Client.serverChain.get(0));
			Message msg = msgin;
			String nextHop;
			boolean destFound = false;
			ServerEntry server;
			EncryptedMessageWrapperMessage wrapper;
			int i = 0;
			do {
				nextHop = null;
				server = chain.get(Constants.numServerHops - i - 1);

				nextHop = server.serverName;// Client.serverChain.get(i).serverName;
				if (server == destination)
					destFound = true;
				if (destFound) {
					Print.debug("{}" + server.serverName + " ->" + nextHop);
					if (!(destination.serverName
							.equalsIgnoreCase(server.serverName) && sendingAESkey)) {
						wrapper = new EncryptedMessageWrapperMessage(
								server.getAsymmetricCipherEngine(), msg,
								Constants.ASYMMETRIC_ENCRYPTION);
						wrapper.setNextHop(nextHop);
						if (useReturnChain)
							wrapper.setConnectionNumber(Constants.USER_DOWNSTREAM);
						else {
							wrapper.setUsePersistentRoute(false);
							wrapper.setConnectionNumber(0);
						}
						msg = wrapper;
					} else {
						((EncryptedMessageWrapperMessage) msg)
								.setNextHop(nextHop);
						((EncryptedMessageWrapperMessage) msg)
								.setConnectionNumber(Constants.USER_DOWNSTREAM);
					}
				}
				i++;
			} while (i < Constants.numServerHops);
			out.writeObject(msg);
			out.flush();
		} catch (Exception e) {
			throw new IOException(
					"Error occured sending message from client to server "
							+ e.getLocalizedMessage());
		}
	}

	public static void main(String[] args) throws ClassNotFoundException,
			NoSuchAlgorithmException, DecryptionException {
		Client client = new Client();
		client.setNickListFrame(new NickListFrame(client));
	}

	public Identity getIdentity() {
		return identity;
	}

	public void setIdentity(Identity identity) {
		this.identity = identity;
	}

	public NickListFrame getNickListFrame() {
		return nickListFrame;
	}

	public void setNickListFrame(NickListFrame nickListFrame) {
		this.nickListFrame = nickListFrame;
	}

	public NickNameList getNickList() {
		return nickList;
	}

	public void setNickList(NickNameList nickList) {
		this.nickList = nickList;
	}

	public PublicIdentity getPublicIdentity() {
		return this.publicIdentity;
	}

	public ObjectOutputStream getOutputStream() {
		return out;
	}

	public Map<String, PublicIdentity> getPublicIdentityMap() {
		return publicIdentityMap;
	}

	/**
	 * Select Identity from previous Identities, or select to create a new ID
	 * 
	 * @return
	 * @throws ClassNotFoundException
	 * @throws IOException
	 * @throws NoSuchAlgorithmException
	 */
	private Identity selectIdentity() throws ClassNotFoundException,
			NoSuchAlgorithmException {
		Identity userID = null;
		NameSaltPairs previous = LocalFileOps.getPreviousUserNamesAndSalt();
		String[] prevUsers = previous.getUserNames();
		String[] optionalUserNames = new String[prevUsers.length + 1];
		optionalUserNames[0] = "New Identity";

		for (int i = 0; i < prevUsers.length; i++) {
			// i is an index into prevUsers optionalUsers is 1 long that
			// prevUsers, with Constant as first element
			optionalUserNames[i + 1] = prevUsers[i];
		}

		String selectedUserName = (String) JOptionPane.showInputDialog(null,
				"Choose an Identity", "Input", JOptionPane.INFORMATION_MESSAGE,
				null, optionalUserNames,
				optionalUserNames[optionalUserNames.length - 1]);
		if (selectedUserName == null) {
			Print.info("User canceled program, goodbye.");
			System.exit(0);
		}
		if (selectedUserName.equals(optionalUserNames[0])) {
			userID = makeNewIdentity(previous);
		} else {
			userID = usePreviousIdentity(previous, selectedUserName);
		}
		if (userID == null) {
			// Tell user they have the wrong password and end program
			LocalFileOps.appendAuditFile("Incorrect password for user "
					+ selectedUserName);
			JOptionPane.showMessageDialog(null,
					"You got the wrong password. Byebye.", "Error",
					JOptionPane.ERROR_MESSAGE);
			System.exit(0);
		}
		return userID;
	}

	/**
	 * Reuse an identity which was created prior to this run. Load the id from
	 * the file system and check its integrity
	 * 
	 * @param previous
	 *            previous salt and name pairs used by the sytem
	 * @param selectedUserName
	 *            username to load
	 * @return the loaded ID
	 */
	public Identity usePreviousIdentity(NameSaltPairs previous,
			String selectedUserName) {
		Identity userID;
		String password;
		// Existing ID
		// prompt password req
		// open file and read id
		Print.debug("Using previous ID. . .");
		KeySalt prevSaltUsed = previous.getUserSalt(selectedUserName);
		if (prevSaltUsed == null) {
			Print.error("User's last salt value is null. Terminating the program");
			System.exit(-1);
		}
		password = promptForPassword();
		SecretKey secret = Keys.createAESKey(password,
				prevSaltUsed.getSalt1Bytes());
		SecretKey secret2 = Keys.createAESKey(password,
				prevSaltUsed.getSalt2Bytes());
		userID = LocalFileOps.getIdentityFromFile(selectedUserName, secret,
				secret2, prevSaltUsed);
		return userID;
	}

	/**
	 * Creates a new user identity and saves it to the file system
	 * 
	 * @param previousPairs
	 * @return
	 */
	public Identity makeNewIdentity(NameSaltPairs previousPairs) {
		Identity userID;
		String password;
		// New ID
		Print.debug("Creating a new ID. . .");
		password = this.promptForPassword();
		if (password == null || password.equalsIgnoreCase("")) {
			Print.error("Cannot use a blank password");
			System.exit(1);
		}
		// Create a new user id
		KeySalt newSalt = new KeySalt();
		JProgressBar progress = this.showProgressPane();
		Print.debug("Creating new id in progress");
		userID = generateNewIdentity(password, newSalt);
		progress.setString("Created Identity");
		progress.setIndeterminate(false);
		progress.setValue(100);
		// add the new user to the previous users of this system
		previousPairs.addPair(userID.getUserName(), newSalt);
		// Update the previous users file
		LocalFileOps.writePreviousUsers(previousPairs);
		return userID;
	}

	/**
	 * Give user a popup asking to provide a password
	 * 
	 * @return
	 */
	private String promptForPassword() {
		JPasswordField passwordField = new JPasswordField();
		Object[] obj = { "Please input a password for this Identity",
				passwordField };
		Object stringArray[] = { "OK", "Cancel" };
		String password = null;
		if (JOptionPane.showOptionDialog(null, obj, "Password",
				JOptionPane.YES_NO_OPTION, JOptionPane.QUESTION_MESSAGE, null,
				stringArray, obj) == JOptionPane.YES_OPTION) {
			password = new String(passwordField.getPassword());
		}
		return password;
	}

	/**
	 * Show a pane with a progress bar that lets the user know something is
	 * happening
	 * 
	 * @return
	 */
	private JProgressBar showProgressPane() {
		final JProgressBar prog = new JProgressBar();
		prog.setIndeterminate(true);
		prog.setString("Creating new identity....");
		prog.setStringPainted(true);
		Thread t = new Thread() {
			@Override
			public void run() {
				Object[] obj = { prog };
				JOptionPane.showMessageDialog(null, obj, "Progress",
						JOptionPane.PLAIN_MESSAGE);
			}
		};
		t.start();
		return prog;
	}

	/**
	 * Generate a new Identity, and save it to the file system encrypted
	 * 
	 * @param password
	 *            - The user's provided password
	 * @param salt
	 *            - KeySalt to use for protecting user's local files
	 * @return The newly created identity
	 */
	private Identity generateNewIdentity(String password, KeySalt salt) {
		Identity n = null;
		try {
			SecretKey encSecret = Keys.createAESKey(password,
					salt.getSalt1Bytes());
			SecretKey macSecret = Keys.createAESKey(password,
					salt.getSalt2Bytes());
			n = new Identity(encSecret, macSecret, salt);
			LocalFileOps.writeUserId(n, encSecret, salt);
		} catch (NoSuchAlgorithmException e) {
			Print.error("Failed to create new Idenity.", e);
			JOptionPane
					.showMessageDialog(
							null,
							"Failed to create a new user identity. Program terminated.",
							"Error", JOptionPane.ERROR_MESSAGE);
			System.exit(0);
		}
		return n;
	}

	/**
	 * Handle encrypted text message with sig.
	 * 
	 * @param textWithSig
	 */
	public void handleRSAEncryptedMessage(RSAEncryptedMessage msg) {
		Print.debug("Get RSAencryptedMessage");
		if (!publicIdentityMap.containsKey(msg.getSrcUserName())
				|| publicIdentityMap.get(msg.getSrcUserName()) == null) {
			String userNameString = PublicIdentity.composeUserName(msg
					.getPublicIdentity().getPublicRSAKey(), msg
					.getPublicIdentity().getVerKey());
			if (msg.getSrcUserName().equals(userNameString)) {
				publicIdentityMap.put(msg.getSrcUserName(),
						msg.getPublicIdentity());
			} else {
				Print.debug("PublicIdentity got compromised!!!!! Under attack!!!!");
				System.exit(-1);
			}
		}

		Print.debug("Decrypting with privatekey");

		TextMessageWithSig textWithSig = SecurityUtils.decRSA(msg, this
				.getIdentity().getPrivateRSAKey());

		if (publicIdentityMap.get(msg.getSrcUserName()) == null) {
			// user is not online
			Print.debug("user" + msg.getSrcUserName() + " is not online.");
		} else {
			boolean verOk = SecurityUtils.ver(textWithSig.getTextMessage(),
					textWithSig.getSignature(), this.getPublicIdentityMap()
							.get(textWithSig.getTextMessage().getSrcUser())
							.getVerKey());
			if (verOk) {
				// signature of text message is verified.
				Print.debug("signature of text message is verified.");
				this.nickListFrame
						.showTextMessage(textWithSig.getTextMessage());
			} else {
				LocalFileOps
						.appendAuditFile("Bad signature of text message from user "
								+ msg.getSrcUserName());
				Print.error("Bad signature of text message");
			}
		}
	}

	public String getPath() {
		String path = null;
		if (this.chain != null) {
			path = this.chain.toString();
		}
		return path;
	}
}
